"""
    Title

    author: wxz
    date: 
    github: https://github.com/xinzwang
"""

import random
from funcs import *


def gauss(loss):
    """
    1. 初始化参数 -> 均值、方差

    2. 根据均值、方差，以高斯分布，获取规模为N的种群
    3. 从种群中采样最好的n个个体，组成精英集合
    4. 根据经营集合，评估均值、方差
    5. 返回 2
    """

    def get_key(elem):
        return elem['y']

    # 1. 初始化参数
    N = 1000  # 种群规模
    n = 300  # 精英个数
    t = 0  # 计数器
    mu = (np.random.random() - 0.5) * 10  # 初始 mu  [-5,5]
    sigma = 10 + np.random.random() * 4
    population = []
    path_x = []
    path_y = []
    path_z = []

    for t in range(10):
        # 2.随机采样种群大小为N的后代种群
        population.clear()
        for z in range(N):
            x = [random.gauss(mu, sigma), random.gauss(mu, sigma)]
            y = loss(x)
            a = {'x': x, 'y': y}
            population.append(a)

        # 3.种群排序    前n个最优样本作为精英集合
        population.sort(key=get_key)

        # 4.使用精英集合评估新的均值和标准差
        mu_t = 0
        sigma_t = 0
        for j in range(n):
            k = (population[j]['x'][0] + population[j]['x'][1]) / 2
            mu_t += k
            sigma_t += (k - mu) ** 2

        mu = mu_t / n
        sigma = sigma_t / n

        path_x.append(population[0]['x'][0])
        path_y.append(population[0]['x'][1])
        path_z.append((population[0]['y']))
        print('t:{}    mu:{}    sigma:{}    x:{} '.format(t, mu, sigma, population[0]))
    return path_x, path_y, path_z


if __name__ == '__main__':
    import matplotlib.pyplot as plt
    f = F_Rosenbrock(n=2)

    loss = f.forward

    xp, yp, zp = gauss(loss)

    fig = plt.figure()
    x = np.arange(-10, 10, 0.1)
    y = np.arange(-10, 10, 0.1)
    x, y = np.meshgrid(x, y)
    z = fun3(x, y)
    ax = fig.add_subplot(111, projection='3d')
    surf = ax.plot_surface(x, y, z, cmap=cm.coolwarm)
    fig.colorbar(surf, shrink=0.5, aspect=5)
    ax.plot(xp, yp, zp, c='r')  # 画点

    plt.xlabel('x')
    plt.ylabel('y')
    plt.show()
